"""
Authors:    Jingjing WU (吴京京) <https://github.com/wj-Mcat>

2020-now @ Jingjing Wu
Licensed under the Apache License, Version 2.0 (the 'License');
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an 'AS IS' BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
from __future__ import annotations

import os
from typing import List, Tuple, Iterable, Dict, Iterator
from torch.utils.data import DataLoader, TensorDataset
from src.config import root_dir
from src.schema import TextClassificationInputExample, SequenceTaggingInputExample


class NLUCsvDataReader:
    """slot filling and intent dataset reader"""
    @staticmethod
    def _read_lines(file: str) -> Iterable[Tuple[str, str, str]]:
        with open(os.path.join(root_dir, file), 'r', encoding='utf-8') as file_handler:

            have_skip_first_line = False
            for line in file_handler:
                if not have_skip_first_line:
                    have_skip_first_line = True
                    continue

                items = line.strip().split(',')
                yield items[0], items[1], items[2]

    def read_intents(self, file: str) -> List[TextClassificationInputExample]:
        """read intent examples from file"""
        examples = []
        for index, (text, _, intent) in enumerate(self._read_lines(file)):
            input_example = TextClassificationInputExample(
                input_id=index,
                raw_text=text,
                category=intent,
                metadata={}
            )
            examples.append(input_example)
        return examples

    def read_slots(self, file: str) -> List[SequenceTaggingInputExample]:
        """read slot examples from file"""
        for index, (text, slots_string, _) in enumerate(self._read_lines(file)):
            input_example = SequenceTaggingInputExample(
                input_id=index,
                raw_text=text,
                metadata={},
                tags=slots_string.split()
            )
            yield input_example

